// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/memory/memory.h>

namespace kumo::pollux::exec {
    namespace test {
        class ScaledScanControllerTestHelper;
    }

    /// Controller used to scales table scan processing based on the query memory
    /// usage.
    class ScaledScanController {
    public:
        /// 'nodePool' is the table scan node pool. 'numDrivers' is number of the
        /// table scan drivers. 'scaleUpMemoryUsageRatio' specifies the memory usage
        /// ratio used to make scan scale up decision.
        ScaledScanController(
            memory::MemoryPool *nodePool,
            uint32_t numDrivers,
            double scaleUpMemoryUsageRatio);

        ~ScaledScanController();

        ScaledScanController(const ScaledScanController &) = delete;

        ScaledScanController(ScaledScanController &&) = delete;

        ScaledScanController &operator=(const ScaledScanController &) = delete;

        ScaledScanController &operator=(ScaledScanController &&) = delete;

        /// Invoked by a scan operator to check if it needs to stop waiting for scan
        /// up to start processing. If so, 'future' is set to a future that will be
        /// ready when the controller decides to start this scan operator processing,
        /// or all the splits from the scan node have been dispatched to finish all
        /// the scan operators. 'driverIcx' is the driver id of the scan operator.
        /// Initially, only the first scan operator at driver index 0 is allowed to
        /// run.
        bool shouldStop(uint32_t driverIdx, ContinueFuture *future);

        /// Invoked by a scan operator to update per-driver memory usage estimation
        /// after finish processing a non-empty split. 'driverIdx' is the driver id of
        /// the scan operator. 'driverMemoryUsage' is the peak memory usage of the
        /// scan operator.
        void updateAndTryScale(uint32_t driverIdx, uint64_t driverMemoryUsage);

        struct Stats {
            uint32_t numRunningDrivers{0};

            std::string toString() const;
        };

        Stats stats() const {
            std::lock_guard<std::mutex> l(lock_);
            return {.numRunningDrivers = numRunningDrivers_};
        }

        /// Invoked by the closed scan operator to close the controller. It returns
        /// true on the first invocation, and otherwise false.
        bool close();

        void testingSetMemoryRatio(double scaleUpMemoryUsageRatio) {
            std::lock_guard<std::mutex> l(lock_);
            *const_cast<double *>(&scaleUpMemoryUsageRatio_) = scaleUpMemoryUsageRatio;
        }

    private:
        // Invoked to check if we can scale up scan processing. If so, call
        // 'scaleUpLocked' for scale up processing.
        void tryScaleLocked(std::optional<ContinuePromise> &driverePromise);

        // Invoked to scale up scan processing by bumping up the number of running
        // scan drivers by one. 'drverPromise' returns the promise to fulfill if the
        // scaled scan driver has been stopped.
        void scaleUpLocked(std::optional<ContinuePromise> &driverePromise);

        // Invoked to check if we need to stop waiting for scale up processing of the
        // specified scan driver. If 'driverIdx' is beyond 'numRunningDrivers_', then
        // a promise is set in 'driverPromises_' and the associated future is returned
        // in 'future'.
        bool shouldStopLocked(
            uint32_t driverIdx,
            kumo::pollux::ContinueFuture *future);

        /// Invoked to update the per-driver memory usage estimation with a new driver
        /// report.
        void updateDriverScanUsageLocked(uint32_t driverIdx, uint64_t memoryUsage);

        memory::MemoryPool *const queryPool_;
        memory::MemoryPool *const nodePool_;
        const uint32_t numDrivers_;
        const double scaleUpMemoryUsageRatio_;

        mutable std::mutex lock_;
        uint32_t numRunningDrivers_{1};

        // The estimated per-driver memory usage of the table scan node.
        uint64_t estimatedDriverUsage_{0};

        // The number of drivers that have reported memory usage.
        uint32_t numDriverReportedUsage_{0};

        // The driver resume promises with one per each driver index.
        std::vector<std::optional<ContinuePromise> > driverPromises_;

        bool closed_{false};

        friend class test::ScaledScanControllerTestHelper;
    };
} // namespace kumo::pollux::exec
